# obj_parser.py

"""
OBJ 文件解析器模块 (多材质支持版)
==================================

此模块提供用于解析 .obj 文件格式的函数，现在支持提取模型的
顶点、材质库(mtllib)、材质使用(usemtl)和面信息。
"""
import os
from typing import List, Dict, Any

# 导入我们新建的 mtl_parser 模块
from . import mtl_parser

# 定义新的返回类型别名，让代码更易读
ObjModelData = Dict[str, Any]

def load_obj(file_path: str) -> ObjModelData:
    """
    加载并解析 .obj 文件，提取顶点、面，并处理材质信息。

    现在支持:
    - v, f: 顶点和面 (与之前相同)
    - mtllib file.mtl: 加载指定的材质库文件。
    - usemtl material_name: 设置当前要使用的材质。

    返回:
        ObjModelData: 一个字典，包含解析后的模型数据，结构如下:
        {
          "vertices": [[x,y,z], ...],
          "materials": {"mat_name": {"Kd": [...]}, ...},
          "face_groups": [{"material": "mat_name", "faces": [[...]]}, ...]
        }
    """
    vertices: List[List[float]] = []
    materials: mtl_parser.MaterialDict = {}
    face_groups: List[Dict[str, Any]] = []

    current_material: str | None = None
    # 用于存储当前材质下的面，直到遇到新的usemtl
    current_faces: List[List[int]] = []

    # 获取.obj文件所在的目录，用于拼接.mtl文件的相对路径
    base_dir = os.path.dirname(file_path)

    try:
        with open(file_path, 'r', encoding='utf-8') as file:
            for line_number, line_content in enumerate(file, 1):
                line = line_content.strip()
                comment_start_index = line.find('#')
                if comment_start_index != -1:
                    line = line[:comment_start_index].strip()
                
                if not line:
                    continue

                parts = line.split()
                if not parts:
                    continue
                
                command = parts[0]

                if command == 'v':
                    # ... (顶点解析逻辑保持不变) ...
                    if len(parts) >= 4:
                        try:
                            vertex = [float(parts[1]), float(parts[2]), float(parts[3])]
                            vertices.append(vertex)
                        except ValueError:
                            print(f"警告 (行 {line_number}): 无法解析顶点坐标: '{line}'")
                
                # 【新增】解析 mtllib 指令
                elif command == 'mtllib':
                    if len(parts) > 1:
                        mtl_file_name = parts[1]
                        # 构建.mtl文件的完整路径
                        mtl_file_path = os.path.join(base_dir, mtl_file_name)
                        print(f"信息: 找到材质库引用，正在加载 '{mtl_file_path}'")
                        # 调用 mtl_parser 来加载材质，并更新到我们的 materials 字典中
                        loaded_mtls = mtl_parser.load_mtl(mtl_file_path)
                        materials.update(loaded_mtls)
                    else:
                        print(f"警告 (行 {line_number}): mtllib 指令缺少文件名。")

                # 【新增】解析 usemtl 指令
                elif command == 'usemtl':
                    if len(parts) > 1:
                        # 当遇到新的 usemtl 指令时，说明上一个材质的面列表已经结束了。
                        # 我们需要将【之前】收集的面（如果存在）保存起来。
                        if current_faces:
                            face_groups.append({
                                "material": current_material,
                                "faces": current_faces
                            })

                        # 开始一个新的材质分组
                        current_material = parts[1]
                        current_faces = [] # 创建一个空列表来收集新材质的面
                    else:
                        print(f"警告 (行 {line_number}): usemtl 指令缺少材质名。")

                elif command == 'f':
                    # ... (面索引解析逻辑基本不变) ...
                    if len(parts) >= 4:
                        face_vertex_indices: List[int] = []
                        valid_face = True
                        for part_str in parts[1:]:
                            vertex_index_str = part_str.split('/')[0]
                            try:
                                vertex_index = int(vertex_index_str)
                                if 1 <= vertex_index <= len(vertices):
                                    face_vertex_indices.append(vertex_index - 1)
                                elif vertex_index < 0:
                                    adjusted_index = len(vertices) + vertex_index
                                    if 0 <= adjusted_index < len(vertices):
                                        face_vertex_indices.append(adjusted_index)
                                    else:
                                        valid_face = False; break
                                else:
                                    valid_face = False; break
                            except ValueError:
                                valid_face = False; break
                        
                        if valid_face and len(face_vertex_indices) >= 3:
                            # 将解析出的面添加到【当前】材质的面列表中
                            current_faces.append(face_vertex_indices)

    except FileNotFoundError:
        print(f"错误: 文件 '{file_path}' 未找到。")
        return {}
    except Exception as e:
        print(f"错误: 解析文件 '{file_path}' 时发生意外错误: {e}")
        return {}

    # 【重要】循环结束后，别忘了保存最后一组正在收集的面
    if current_faces:
        face_groups.append({
            "material": current_material,
            "faces": current_faces
        })

    # 封装成我们设计好的新数据结构并返回
    return {
        "vertices": vertices,
        "materials": materials,
        "face_groups": face_groups
    }

# --- 模块独立测试代码 ---
if __name__ == "__main__":
    print("--- 开始独立测试 obj_parser.py (多材质版) ---")

    # 1. 创建临时的 .mtl 文件
    test_mtl_content = """
    newmtl Red
    Kd 1.0 0.0 0.0
    newmtl Blue
    Kd 0.0 0.0 1.0
    """
    test_mtl_path = "test_multi.mtl"
    with open(test_mtl_path, 'w', encoding='utf-8') as f:
        f.write(test_mtl_content)
    print(f"已创建测试MTL文件: {test_mtl_path}")

    # 2. 创建临时的 .obj 文件，引用上面的 .mtl
    test_obj_content = f"""
# 多材质测试 OBJ
mtllib {test_mtl_path}

v 0 0 0
v 1 0 0
v 1 1 0
v 0 1 0

# 第一个物体，红色
usemtl Red
f 1 2 3

# 第二个物体，蓝色
usemtl Blue
f 1 3 4

# 第三个物体，没有材质
usemtl off
f 2 4 3
"""
    test_obj_path = "test_multi.obj"
    with open(test_obj_path, 'w', encoding='utf-8') as f:
        f.write(test_obj_content)
    print(f"已创建测试OBJ文件: {test_obj_path}")
    
    # 3. 调用新的 load_obj
    model_data = load_obj(test_obj_path)
    
    # 4. 打印并验证结果
    print("\n--- 解析结果 ---")
    # 为了美观，我们不直接打印整个大字典，而是分部分打印
    print("Vertices:", model_data.get("vertices"))
    print("Materials:", model_data.get("materials"))
    print("Face Groups:")
    import json
    print(json.dumps(model_data.get("face_groups"), indent=2))

    print("\n--- 结果验证 ---")
    assert len(model_data["vertices"]) == 4, "顶点数量不正确"
    assert "Red" in model_data["materials"], "红色材质未加载"
    assert model_data["materials"]["Red"]["Kd"] == [1.0, 0.0, 0.0], "红色材质颜色不正确"
    assert len(model_data["face_groups"]) == 3, "面分组数量应为3"
    
    red_group = model_data["face_groups"][0]
    assert red_group["material"] == "Red", "第一组材质应为Red"
    assert len(red_group["faces"]) == 1, "Red组应有1个面"
    assert red_group["faces"][0] == [0, 1, 2], "Red组的面索引不正确"

    off_group = model_data["face_groups"][2]
    assert off_group["material"] == "off", "第三组材质应为off (usemtl off)" # OBJ规范中 usemtl off 表示关闭材质

    print("所有断言测试通过！")

    # 5. 清理
    os.remove(test_mtl_path)
    os.remove(test_obj_path)
    print("\n已清理测试文件。")
    print("--- obj_parser.py (多材质版) 测试结束 ---")